# 多臂老虎机问题

import numpy as np
import matplotlib.pyplot as plt

class BernoulliBandit:
    def __init__(self, k):
        self.probs = np.random.uniform(size=k)  # 随机生成k个0-1之间的数
        self.best_idx = np.argmax(self.probs)  # 最大的数的索引
        self.best_prob = self.probs[self.best_idx]  # 最大的数
        self.K = k

    def step(self, k):
        # 当玩家选择了k好栏杆后，根据概率返回1（获奖）或者0（未获奖）
        if np.random.rand() < self.probs[k]:
            return 1
        else:
            return 0

class Solver:
    def __init__(self, bandit):
        self.bandit = bandit
        self.counts = np.zeros(self.bandit.K)  # 每根拉杆的尝试次数
        self.regret = 0.  # 当前步的累积懊悔
        self.actions = []  # 维护一个列表，记录每一步的动作
        
        self.regrets = []  # 维护一个列表，记录每一步的累积懊悔

    def update_regret(self, k):
        # 计算累积懊悔并保存
        self.regret += self.bandit.best_prob - self.bandit.probs[k]
        self.regrets.append(self.regret)

    def run_one_step(self):
        # 返回当前动作选择哪一根拉杆，由每个具体的策略实现
        raise NotImplementedError

    def run(self, num_steps):
        # 运行一定次数，runs是一个列表，记录每一步的动作
        for _ in range(num_steps):
            k = self.run_one_step()
            self.counts[k] += 1
            self.actions.append(k)
            self.update_regret(k)



def plot_results(solvers, solver_names):
    """生成累积懊悔随时间变化的图像。输入solvers是一个列表，列表中的每个元素是一种特定的策略。
    而solver_names也是一个列表，存储每个策略的名称"""
    for idx, solver in enumerate(solvers):
        time_list = range(len(solver.regrets)) 
        plt.plot(time_list, solver.regrets, label=solver_names[idx])
    plt.xlabel('Time steps')
    plt.ylabel('Cumulative regrets')
    plt.title("%d-armed bandit" % solvers[0].bandit.K)
    plt.legend()
    plt.show()

k = 10
bandit_10_arm = BernoulliBandit(k)
if __name__ == '__main__':
    np.random.seed(1)
    print("随机生成了一个%d臂伯努利老虎机" % k)
    print("获奖概率最大的拉杆为%d号，其获奖概率为%.4f" %(bandit_10_arm.best_idx, bandit_10_arm.best_prob))
